package cn.com.basic.framework.aliyun.odps.tunnel;

import cn.com.basic.framework.aliyun.odps.OdpsManager;
import cn.com.basic.framework.aliyun.odps.error.OdpsError;
import cn.com.basic.framework.aliyun.odps.utils.OdpsUtil;
import cn.com.basic.framework.exception.exceptions.SysException;
import cn.com.basic.framework.utils.RetryUtil;
import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.lang.Assert;
import com.aliyun.odps.OdpsException;
import com.aliyun.odps.Table;
import com.aliyun.odps.data.Record;
import com.aliyun.odps.data.RecordReader;
import com.aliyun.odps.tunnel.TableTunnel;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Callable;

/**
 * tunnel 读取 odps 数据
 *
 * @author zhuxj
 */
public class DownloadTunnel {
    private static Logger logger = LoggerFactory.getLogger(DownloadTunnel.class);
    private OdpsManager manager;
    private String tunnelServer;

    public DownloadTunnel(OdpsManager manager) {
        this.manager = manager;
    }

    /**
     * 普通表读取
     *
     * @param projectName 项目空间名称
     * @param tableName   表名称
     */
    public List<Record> download(String projectName, String tableName) {
        return download(projectName, tableName, null);
    }

    /**
     * 分区表读取
     *
     * @param projectName 项目名称
     * @param tableName   表名称
     * @param partition   分区值
     */
    public List<Record> download(String projectName, String tableName, List<String> partition) {
        Assert.isTrue(StringUtils.isNoneBlank(projectName), "project is not presented");
        Assert.isTrue(StringUtils.isNoneBlank(tableName), "table is not presented");
        Assert.isTrue(checkConfig());

        Table table = manager.getTable(projectName, tableName);
        if (table.isVirtualView()) {
            throw SysException.newException(OdpsError.TUNNEL_ERROR, String.format("目标读取的表(视图)是%s,Tunnel通道不支持读取虚拟视图", tableName));
        }


        boolean isPartitionTable;
        try {
            isPartitionTable = table.isPartitioned();

            if (isPartitionTable && CollectionUtil.isEmpty(partition))
                throw SysException.newException(OdpsError.PARTITION_ERROR, "检查表%s为分区表，但是到未监测到partition参数，分区表读取需要指定分区，格式形如:[\"pt=hello,ds=hangzhou\",\"pt=world,ds=hangzhou\"],若要全表读取可填'*'");

            if (isPartitionTable)
                OdpsUtil.checkPartition(table, partition);

        } catch (OdpsException e) {
            logger.error(e.getMessage(), e);
            throw SysException.newException(OdpsError.PARTITION_ERROR, e);
        }


        TableTunnel.DownloadSession session;
        if (isPartitionTable) {
            session = OdpsUtil.createMasterSessionForPartitionedTable(this.manager.getOdps(),
                    tunnelServer, projectName, tableName, partition.get(0));
        } else {
            session = OdpsUtil.createMasterSessionForNonPartitionedTable(this.manager.getOdps(),
                    tunnelServer, projectName, tableName);

        }


        long start = 0;
        long count = session.getRecordCount();

        if (count > 0) {
            logger.info(String.format(
                    "Begin to read ODPS table:%s, partition:%s, startIndex:%s, count:%s.",
                    tableName, partition, start, count));
        } else if (count == 0) {
            logger.warn(String.format("源头表:%s 的分区:%s 没有内容可抽取, 请您知晓.",
                    tableName, partition));
            return new ArrayList<>();
        } else {
            throw SysException.newException(OdpsError.READ_DATA_ERROR,
                    String.format("源头表:%s 的分区:%s  读取行数为负数, 请联系 ODPS 管理员查看表状态!",
                            tableName, partition));
        }

        RecordReader recordReader = OdpsUtil.getRecordReader(session, start, count, false);
        List<Record> results = new ArrayList<>();
        while (true) {
            Record record = null;
            try {
                record = RetryUtil.executeWithRetry(new Callable<Record>() {
                    @Override
                    public Record call() throws Exception {
                        return recordReader.read();

                    }
                }, 10, 2000, false);
            } catch (Exception e) {
                logger.error("读取重试失败，失败原因:{}", e.getLocalizedMessage());
                logger.error(e.getMessage(), e);
                throw SysException.newException(OdpsError.READ_DATA_ERROR, e);
            }
            start++;
            count--;
            if (record != null) {
                results.add(record);
            } else {
                break;
            }
        }

        try {
            recordReader.close();
        } catch (Exception e) {
            logger.warn("recordReader close exception", e);
        }

        return results;
    }

    public void setTunnelServer(String tunnelServer) {
        this.tunnelServer = tunnelServer;
    }

    private boolean checkConfig() {
        if (this.manager == null) {
            throw SysException.newException(OdpsError.MANAGER_IS_NULL, "尚未配置OdpsManager");
        }
        if (this.manager.getOdps() == null) {
            throw SysException.newException(OdpsError.MANAGER_IS_NULL, "管理器尚未配置ODPS客户端");
        }
        return true;
    }
}
